# Author：gyx
# Usage: build uav with robot arm dynamics model
# Last change: 2024-7-20

import numpy as np
from common_function import dcm_from_euler, rotation_matrix,angle_body2world,Skew_symmetric_matrix,cross_product
from init_param import UAVParameters

class UAVModel:
    def __init__(self, params):
        self.params = params
        self.w_current = params.ModelInit_Rads
        self.xi=np.concatenate(params.ModelInit_PosE,params.ModelInit_AngEuler,params.ModelInit_Jointq)

    def get_angular_velocity_derivative(self, M, J, omega):
        return np.linalg.inv(J) @ (M - np.cross(omega, J @ omega))

    def get_linear_velocity_derivative(self, F, m):
        return np.clip(F / m, -0.4, 0.4)

    def get_state_derivate(self,state,F):
        posB, velB, angEuler, rateB,posEndB,anglejoint = state
        B11=(self.params.ModelParam_uavMass+self.params.linkmass1+self.params.linkmass2)*np.eye(3)
        R=rotation_matrix(angEuler)
        Q=np.linalg.inv(angle_body2world(angEuler))
        S1=Skew_symmetric_matrix(posEndB(0))
        S2=Skew_symmetric_matrix(posEndB(1))
        Rl1=rotation_matrix(anglejoint(0))
        Rl2=rotation_matrix(anglejoint(1))
        Il1=np.array([[1/3*self.params.linkmass1*self.params.linklength1^2,0,0],
                     [0,1/4*self.params.linkmass1*self.params.linklength1^2,0],
                     [0,0,1/3*self.params.linkmass1*self.params.linklength1^2]])
        Il2=np.array([[1/3*self.params.linkmass2*self.params.linklength2^2,0,0],
                     [0,1/4*self.params.linkmass2*self.params.linklength2^2,0],
                     [0,0,1/3*self.params.linkmass2*self.params.linklength2^2]])
        B22=Q.T@self.params.ModelParam_uavJ@Q\
            +(self.params.linkmass1*Q.T@S1).T@S1@Q+Q.T@Rl1@Il1@Rl1.T@Q\
            +(self.params.linkmass2*Q.T@S2).T@S2@Q+Q.T@Rl2@Il2@Rl2.T@Q
        e2=np.array([[0],
                       [1],
                       [0]])
        Jpl1,Jpl2=np.array([[0,0],
                       [1,1],
                       [0,0]])
        ol1=cross_product([0,1,0],posEndB(0)-posB)
        ol2=cross_product([0,1,0],posEndB(1)-posEndB(0))
        Jol1=np.array([[ol1(0),0],
                       [ol1(1),0],
                       [ol1(2),0]])
        Jol2=np.array([[ol1(0),ol2(0)],
                       [ol1(1),ol2(1)],
                       [ol1(2),ol2(2)]])
        B33=self.params.linkmass1@Jpl1.T@Jpl1+Jol1@Rl1@Il1@Rl1.T@Jol1\
            +self.params.linkmass2@Jpl2.T@Jpl2+Jol2@Rl2@Il2@Rl2.T@Jol2
        B12=-(self.params.linkmass1@R@S1@Q+self.params.linkmass2@R@S2@Q)
        B13=self.params.linkmass1@R@Jpl1+self.params.linkmass2@R@Jpl2
        B23=Q.T@Rl1@Il1@Rl1.T@Jol1-self.params.linkmass1@Q.T@S1.T@Jpl1\
            +Q.T@Rl2@Il2@Rl2.T@Jol2-self.params.linkmass2@Q.T@S2.T@Jpl2
        top_row = np.hstack((B11, B12, B13))
        middle_row = np.hstack((B12.T, B22, B23))
        bottom_row = np.hstack((B13.T, B23.T, B33))
        B = np.vstack((top_row, middle_row, bottom_row))
    def Obtain_force_torque(self, w, R, Cm, Ct, Vb, Cd, wb, Cdm, Jrp):
        M_rctcm = np.array([
            [-np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct,
             -np.sin(np.pi / 4) * R * Ct],  # Roll torque
            [np.sin(np.pi / 4) * R * Ct, -np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct,
             -np.sin(np.pi / 4) * R * Ct],  # Pitch torque
            [Cm, Cm, -Cm, -Cm]  # Yaw torque
        ])

        Mp = np.dot(M_rctcm, w ** 2)  # Torque

        Fp = np.array([0, 0, -np.sum(Ct * (w ** 2))])  # Thrust
        Fd = -Cd * Vb * np.abs(Vb) * 0.5  # Aerodynamic force
        Md = -Cdm * wb * np.abs(wb)  # Aerodynamic moment

        # Gyro moment
        Ga = np.zeros(3)
        Ga[0] = Jrp * wb[1] * (w[0] + w[1] - w[2] - w[3])
        Ga[1] = Jrp * wb[0] * (-w[0] - w[1] + w[2] + w[3])

        return Fp, Fd, Mp, Md, Ga

    def motor_dynamics(self, w_current, w_target, T, dt):
        dw = (w_target - w_current) / T
        return w_current + dw * dt

    def state_derivative(self, state, w_target, dt):
        velB, angEuler, rateB = state

        m = self.params.ModelParam_uavMass
        J = self.params.ModelParam_uavJ
        g = self.params.ModelParam_envGravityAcc
        uavType = self.params.ModelParam_uavType
        R = self.params.ModelParam_uavR
        Cm = self.params.ModelParam_rotorCm
        Ct = self.params.ModelParam_rotorCt
        Cd = self.params.ModelParam_uavCd
        Cdm = self.params.ModelParam_uavCCm
        Jrp = self.params.ModelParam_motorJm
        T = self.params.ModelParam_motorT

        self.w_current = self.motor_dynamics(self.w_current, w_target, T, dt)
        Fp, Fd, Mp, Md, Ga = self.Obtain_force_torque(self.w_current, R, Cm, Ct, velB, Cd, rateB, Cdm, Jrp)
        dcm = dcm_from_euler(angEuler)
        total_force = Fp + Fd + np.array([0, 0, m * g]).dot(dcm)
        total_torque = Mp + Md + Ga
        velB_dot = self.get_linear_velocity_derivative(total_force, m)
        rateB_dot = self.get_angular_velocity_derivative(total_torque, J, rateB)

        return velB_dot, rateB_dot

    def update_state(self, state, w_target, dt):
        velB, angEuler, rateB = state

        # Runge-Kutta 4th Order Method
        k1_velB_dot, k1_rateB_dot = self.state_derivative(state, w_target, dt)
        k2_velB_dot, k2_rateB_dot = self.state_derivative(
            (velB + 0.5 * k1_velB_dot * dt, angEuler, rateB + 0.5 * k1_rateB_dot * dt), w_target, dt)
        k3_velB_dot, k3_rateB_dot = self.state_derivative(
            (velB + 0.5 * k2_velB_dot * dt, angEuler, rateB + 0.5 * k2_rateB_dot * dt), w_target, dt)
        k4_velB_dot, k4_rateB_dot = self.state_derivative(
            (velB + k3_velB_dot * dt, angEuler, rateB + k3_rateB_dot * dt), w_target, dt)

        velB = velB + (k1_velB_dot + 2 * k2_velB_dot + 2 * k3_velB_dot + k4_velB_dot) * dt / 6
        rateB = rateB + (k1_rateB_dot + 2 * k2_rateB_dot + 2 * k3_rateB_dot + k4_rateB_dot) * dt / 6

        return velB, angEuler, rateB
